#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
@author: xjx
@time: 2024/3/18 10:05 
@file: load_analysis.py
@project: nanchangproject
@describe: 用于载荷分析的模型
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 仅与车辆参数相关的变量
ig = 4.11  # 主减速比
Torque_max = 450  # 最大扭矩 Nm
A_wind = 2.29 * 2.16  # 车辆迎风面积 m^2

# 可能随状态变化的变量
gravity = 9.81  # 重力加速度 m/s^2
air_density = 1.225  # 空气密度
air_coefficient = 0.3  # 空气阻力系数
road_coefficient = 0.015  # 滚动阻力系数
radius = 0.34284  # 轮胎滚动半径 m

# EKF相关变量
std_Torque = 5 / Torque_max  # 扭矩噪声标准差
std_v = 0.1  # 速度噪声标准差
std_inverse_m = 100  # 载荷噪声标准差(非倒数)
std_alpha = 0.03  # 坡度噪声标准差


def transmission_efficiency(u_k):  # 确定当前传动效率
    Torque = [0.0653, 0.0816, 0.1, 0.1490, 0.1939, 0.2490, 0.3, 0.3816, 0.4837, 0.4959, 0.6061, 0.6224, 0.7735, 0.7939,
              1]  # 扭矩百分比 共15档
    eta = [80.2105, 84.3158, 87.3684, 90.9474, 92.3158, 93.4737, 94.2105, 94.7368, 95.3684, 95.4737, 95.7895, 95.8947,
           96.3158, 96.4211, 96.3158]  # 传动效率
    closest_Torque = min(Torque, key=lambda x: abs(x - u_k))  # 查找最接近的扭矩
    eta_position = Torque.index(closest_Torque)  # 确定传动效率位置
    eta = eta[eta_position]  # 确定传动效率
    return eta


def determine_gear(engine_rpm, v_kmph):  # 确定当前挡位和传动比
    gear_ratios = [5.701, 5.797, 3.434, 1.862, 1.29, 1, 0.759]  # 变速器档速比 R123456
    v_mpm = v_kmph * 1000 / 60  # 车速单位转换
    if v_mpm != 0:
        actual_it = (engine_rpm / v_mpm)  # 计算实际传动比
    else:
        actual_it = 0
    closest_gear = min(gear_ratios, key=lambda x: abs(x - actual_it))  # 查找最接近的已知档速比
    gear_position = gear_ratios.index(closest_gear)  # 确定挡位
    it = gear_ratios[gear_position]  # 确定档速比
    # it = actual_it # 用实际传动比代替
    return gear_position, it


def forward_euler_discretization(x_k, u_k, v_mps):  # 前向欧拉离散化
    """
    较好的结果：速度不考虑噪声，扭矩 载荷 坡度考虑噪声
    std_alpha对结果影响最显著
    """
    # Torque_hat = max(0,min(1,u_k*Torque_max)) # 不考虑噪声的扭矩
    Torque_hat = max(0, min(1, u_k * Torque_max + omega_Torque))  # 考虑噪声的扭矩
    v_hat = max(0, min(90 / 3.6, v_mps))  # 不考虑噪声的速度
    # v_hat = max(0,min(90/3.6,v_mps+omega_v)) # 考虑噪声的速度
    # v_hat = x_k[0]+omega_v # 考虑噪声的非真实速度
    # inverse_m_hat = max(1/7495,min(1/2600,x_k[1])) # 不考虑噪声的载荷
    inverse_m_hat = max(1 / 6100, min(1 / 2800, 1 / (1 / x_k[1] + omega_inverse_m)))  # 考虑噪声的载荷
    # alpha_hat = max(-15,min(15,x_k[2])) # 不考虑噪声的坡度
    alpha_hat = max(-15, min(15, x_k[2] + omega_alpha))  # 考虑噪声的坡度
    x_kplus1 = [0.0, 0.0, 0.0]
    x_kplus1[0] = max(0, min(90 / 3.6, v_hat + Ts / delta * (it * ig * eta / radius * Torque_hat * inverse_m_hat \
                                                             - air_density * air_coefficient * A_wind / 2 * inverse_m_hat * v_hat ** 2 \
                                                             - road_coefficient * gravity * np.cos(
                alpha_hat) - gravity * np.sin(alpha_hat))))
    x_kplus1[1] = inverse_m_hat
    x_kplus1[2] = alpha_hat
    return x_kplus1


def extended_average_filter(x_k, u_k, P, x_pred, v_mps):  # 扩展均值滤波
    """
    较好的结果：更新x用真实速度，不用预测结果。P要随着W和Q更新。
    用预测的z会导致算法波动明显
    """
    Q = np.diag([std_v ** 2, std_inverse_m ** 2, std_alpha ** 2])  # 状态噪声协方差矩阵
    R = np.diag([std_v ** 2])  # 观测噪声协方差矩阵
    H = np.array([[1, 0, 0]])
    V = np.array([[1]])

    a11 = 1 - Ts * air_density * air_coefficient * A_wind / delta * x_k[1] * x_k[0]
    a12 = it * ig * eta / radius / delta * Ts * u_k * Torque_max - air_density * air_coefficient * A_wind / 2 / delta * Ts * \
          x_k[0] ** 2
    a13 = Ts * road_coefficient * gravity / delta * np.sin(x_k[2]) - Ts * gravity / delta * np.cos(x_k[2])
    A = np.array([[a11, a12, a13], [0, 1, 0], [0, 0, 1]], dtype=float)
    w11 = 1 - Ts * air_density * air_coefficient * A_wind / delta * x_k[1] * x_k[0]
    w12 = it * ig * eta / radius / delta * Ts * u_k * Torque_max
    w13 = -Ts * gravity / delta
    W = np.array([[w11, w12, w13], [0, 1, 0], [0, 0, 1]], dtype=float)

    z = H @ x_pred + V @ np.array([omega_v])  # 预测下一观测值
    z[0] = max(0, min(90 / 3.6, z[0]))  # 观测值限幅
    P_pred = A @ P @ A.T + W @ Q @ W.T
    # P_pred = A@P@A.T+Q
    K = P_pred @ H.T @ np.linalg.inv(H @ P_pred @ H.T + R)
    # x = x_pred+K@(z-H@x_pred)
    x = x_pred + K @ (v_mps - H @ x_pred)  # 用真实车速替代z
    P = (np.eye(3) - K @ H) @ P_pred
    return x, P


def main_loop_actual(time, engine_rpm, v_kmph, v_mps, Torque):
    global Ts, eta, it, delta, omega_Torque, omega_v, omega_inverse_m, omega_alpha
    Ts = 10  # 时间步长 s
    omega_Torque = 0
    omega_v = 0
    omega_inverse_m = 0
    omega_alpha = 0
    u0 = 350  # 初始控制变量 扭矩
    x0 = [40 / 3.6, 1 / 4500, 0]  # 初始状态变量 [速度,1/载荷,坡度]
    P0 = np.eye(3)  # 初始状态估计协方差矩阵
    x_k = x0
    P_k = P0
    x_k_dict = {}  # 用于存储每步x_k的字典

    try:
        for i in range(len(engine_rpm)):
            if i != len(engine_rpm) - 1:
                Ts = time[i + 1] - time[i]
            omega_Torque = np.random.normal(0, std_Torque)  # 扭矩噪声
            omega_v = np.random.normal(0, std_v)  # 速度噪声
            omega_inverse_m = np.random.normal(0, std_inverse_m)  # 载荷噪声
            omega_alpha = np.random.normal(0, std_alpha)  # 坡度噪声
            eta = transmission_efficiency(Torque[i])
            current_gear, it = determine_gear(engine_rpm[i], v_kmph[i])
            delta = 1 + 0.04 * it + 0.0025 * it * ig  # 旋转质量换算系数
            x_pred = forward_euler_discretization(x_k, Torque[i], v_mps[i])  # 预测下一状态
            x_k, P_k = extended_average_filter(x_k, Torque[i], P_k, x_pred, v_mps[i])
            x_k_dict[i] = [x_pred[0], 1 / x_pred[1], x_pred[2]]  # 保存当前x_k
    except Exception as e:
        print("发生异常：", str(e))
        pass
    return x_k_dict, P_k  # 返回状态变量


def plot_data(data_dict):
    # time_steps = list(data_dict.keys())
    start = 0
    length = 500
    time_steps = list(state_dict.keys())[start:start + length]  # 仅绘制前边的数据
    reciprocals = [data_dict[i][1] for i in time_steps]
    plt.figure(figsize=(10, 6))
    plt.plot(time_steps, reciprocals, label='Mass')
    plt.xlabel('Time Steps')
    plt.ylabel('Values')
    plt.title('Values over Time')
    plt.legend()
    plt.grid(True)
    plt.show()


def read_csv_actual(csv_file):  # 读取csv文件
    time = []
    engine_rpm = []
    v_kmph = []
    v_mps = []
    Torque = []
    df = pd.read_csv(csv_file)
    time_column = df['时间戳']
    engine_rpm_column = df['发动机转速']
    speed_column = df['车速']
    engine_torque_column = df['发动机净输出扭矩或发动机实际扭矩/指示扭矩']
    friction_torque_column = df['摩擦扭矩']
    for i in range(len(engine_rpm_column)):
        try:
            rpm_value = float(engine_rpm_column[i])
            speed_value = float(speed_column[i])
            engine_torque_value = float(engine_torque_column[i])
            friction_torque_value = float(friction_torque_column[i])  # 处理摩擦扭矩列
            if rpm_value != 0 and speed_value != 0:
                time.append(time_column[i])
                engine_rpm.append(rpm_value)
                v_kmph.append(speed_value)
                speed_value_mps = speed_value / 3.6  # 车速转换为 m/s
                v_mps.append(speed_value_mps)
                total_torque = (engine_torque_value + friction_torque_value) / 100  # 相加得到最终发动机扭矩
                Torque.append(total_torque)
        except Exception as e:
            print(e)
            continue
    return time, engine_rpm, v_kmph, v_mps, Torque


def load_analysis_function(file_name: str):
    time, engine_rpm, v_kmph, v_mps, Torque = read_csv_actual(file_name)  # 读取csv文件
    state_dict, final_covariance = main_loop_actual(time, engine_rpm, v_kmph, v_mps, Torque) # 运行主循环
    plot_data(state_dict)


if __name__ == '__main__':
    my_file_name = 'vehicle_data_combined_1205-667(1)-1.csv'
    # time, engine_rpm, v_kmph, v_mps, Torque = read_csv_actual(file_name)  # 读取csv文件
    # state_dict, final_covariance = main_loop_actual(time, engine_rpm, v_kmph, v_mps, Torque)  # 运行主循环
    # # print(state_dict)
    # plot_data(state_dict)
    load_analysis_function(file_name=my_file_name)
